import os
import re
import fitz  # PyMuPDF
from pathlib import Path
from PIL import Image
import io

def extract_images_from_pdf(pdf_path, output_folder):
    """
    从PDF中提取图片并保存到指定文件夹，只保留大于20KB的图片
    如果单页图片超过3张，则导出整页的非文字内容为一张图片
    :param pdf_path: PDF文件路径
    :param output_folder: 输出文件夹路径
    :return: (成功提取的图片数量, 跳过的图片数量, 导出的整页数量)
    """
    doc = fitz.open(pdf_path)
    saved_count = 0
    skipped_count = 0
    fullpage_count = 0
    
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        image_list = page.get_images(full=True)
        
        # 如果单页图片超过3张，导出整页的非文字内容
        if len(image_list) > 3:
            # 创建一个只包含非文字内容的图像
            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2), 
                                 clip=page.rect,
                                 alpha=False,
                                 colorspace=fitz.csRGB,
                                 annots=False)  # 不包含注释
            
            # 转换为PIL图像以便进一步处理
            img = Image.open(io.BytesIO(pix.tobytes()))
            
            # 保存合并后的图像
            fullpage_path = os.path.join(output_folder, f"fullpage_{page_num+1}.png")
            img.save(fullpage_path)
            fullpage_count += 1
            continue
        
        # 正常提取单张图片
        for img_index, img in enumerate(image_list):
            xref = img[0]
            base_image = doc.extract_image(xref)
            image_bytes = base_image["image"]
            image_ext = base_image["ext"]
            
            # 检查图片大小是否大于20KB
            if len(image_bytes) <= 20 * 1024:  # 20KB
                skipped_count += 1
                continue
            
            # 生成图片文件名
            image_filename = f"image{page_num+1}_{img_index+1}.{image_ext}"
            image_path = os.path.join(output_folder, image_filename)
            
            # 保存图片
            with open(image_path, "wb") as image_file:
                image_file.write(image_bytes)
            saved_count += 1
    
    return saved_count, skipped_count, fullpage_count

def get_largest_font_text(page):
    """
    获取页面中字号最大的文本块
    :param page: PDF页面对象
    :return: 最大字号的文本内容
    """
    blocks = page.get_text("dict")["blocks"]
    max_size = 0
    largest_text = ""
    
    for block in blocks:
        if "lines" in block:
            for line in block["lines"]:
                for span in line["spans"]:
                    if span["size"] > max_size:
                        max_size = span["size"]
                        largest_text = span["text"]
                    elif span["size"] == max_size:
                        largest_text += " " + span["text"]
    
    return largest_text.strip()

def get_article_title(pdf_path):
    """
    从PDF中提取文章标题（获取第一页字号最大的文本）
    :param pdf_path: PDF文件路径
    :return: 文章标题（清理后的）
    """
    try:
        doc = fitz.open(pdf_path)
        first_page = doc.load_page(0)
        
        # 获取字号最大的文本作为标题
        title = get_largest_font_text(first_page)
        
        # 清理标题中的特殊字符和换行
        title = re.sub(r'[\\/*?:"<>|\n\r]', "", title).strip()
        
        # 如果标题过长，可能是获取了多行，取第一行
        if len(title) > 200:
            title = title.split(". ")[0] if ". " in title else title[:200]
        
        return title[:200]  # 限制标题长度
    
    except Exception as e:
        print(f"Error extracting title from {pdf_path}: {str(e)}")
        return ""

def process_pdfs(input_folder, output_root):
    """
    处理输入文件夹中的所有PDF文件
    :param input_folder: 输入文件夹路径
    :param output_root: 输出根文件夹路径
    """
    # 确保输出文件夹存在
    Path(output_root).mkdir(parents=True, exist_ok=True)
    
    # 遍历输入文件夹中的所有PDF文件
    for filename in os.listdir(input_folder):
        if filename.lower().endswith(".pdf"):
            pdf_path = os.path.join(input_folder, filename)
            
            try:
                # 获取PDF文件名（不带扩展名）和文章标题
                pdf_name = os.path.splitext(filename)[0]
                article_title = get_article_title(pdf_path)
                
                # 创建组合文件夹名
                if article_title:
                    folder_name = f"{pdf_name} - {article_title}"
                else:
                    folder_name = pdf_name
                
                # 清理文件夹名中的非法字符
                folder_name = re.sub(r'[\\/*?:"<>|]', "", folder_name).strip()
                article_folder = os.path.join(output_root, folder_name)
                Path(article_folder).mkdir(parents=True, exist_ok=True)
                
                # 提取图片（只保留大于20KB的）
                saved_count, skipped_count, fullpage_count = extract_images_from_pdf(pdf_path, article_folder)
                
                print(f"Processed: {filename}")
                print(f"  Title: {article_title}")
                print(f"  Output folder: {folder_name}")
                print(f"  Saved {saved_count} images, skipped {skipped_count} small images (<20KB)")
                print(f"  Exported {fullpage_count} full pages (when >3 images/page)")
                print(f"  Output path: {article_folder}\n")
                
            except Exception as e:
                print(f"Error processing {filename}: {str(e)}\n")

if __name__ == "__main__":
    # 配置输入输出文件夹路径
    INPUT_FOLDER = "papers"  # 替换为您的输入文件夹路径
    OUTPUT_FOLDER = "output_images"  # 替换为您的输出文件夹路径
    
    print("Starting PDF image extraction with enhanced features...")
    print(f"Input folder: {INPUT_FOLDER}")
    print(f"Output folder: {OUTPUT_FOLDER}")
    print("Features:")
    print(" - Extract largest font text as title")
    print(" - Save only images >20KB")
    print(" - Export full page non-text content as image when >3 images on a page\n")
    
    process_pdfs(INPUT_FOLDER, OUTPUT_FOLDER)
    
    print("Processing completed.")